package com.xiaoyu.cloud.core.loadBalancer;


import com.xiaoyu.cloud.core.constants.ComConst;
import com.xiaoyu.common.core.currentUser.TransmittableThreadLocalRequestHeaders;
import org.apache.commons.lang3.StringUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.ObjectProvider;
import org.springframework.cloud.client.ServiceInstance;
import org.springframework.cloud.client.loadbalancer.*;
import org.springframework.cloud.loadbalancer.core.NoopServiceInstanceListSupplier;
import org.springframework.cloud.loadbalancer.core.ReactorServiceInstanceLoadBalancer;
import org.springframework.cloud.loadbalancer.core.ServiceInstanceListSupplier;
import org.springframework.http.HttpHeaders;
import reactor.core.publisher.Mono;

import java.util.ArrayList;
import java.util.List;
import java.util.Objects;
import java.util.Random;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.stream.Collectors;


/**
 * Description:
 * 自定义灰度
 * 通过给请求头添加Version 与 Service Instance 元数据属性进行对比
 */
public class VersionGrayLoadBalancer implements ReactorServiceInstanceLoadBalancer {
    private static final Logger log = LoggerFactory.getLogger(VersionGrayLoadBalancer.class);
    //提供服务列表
    private final ObjectProvider<ServiceInstanceListSupplier> serviceInstanceListSupplierProvider;
    //服务id
    private final String serviceId;

    private final AtomicInteger position;

    public VersionGrayLoadBalancer(ObjectProvider<ServiceInstanceListSupplier> serviceInstanceListSupplierProvider, String serviceId) {
        this(serviceInstanceListSupplierProvider, serviceId, new Random().nextInt(1000));
    }

    public VersionGrayLoadBalancer(ObjectProvider<ServiceInstanceListSupplier> serviceInstanceListSupplierProvider,
                                   String serviceId, int seedPosition) {
        this.serviceId = serviceId;
        this.serviceInstanceListSupplierProvider = serviceInstanceListSupplierProvider;
        this.position = new AtomicInteger(seedPosition);
    }


    @Override
    public Mono<Response<ServiceInstance>> choose(Request request) {
        ServiceInstanceListSupplier supplier = this.serviceInstanceListSupplierProvider.getIfAvailable(NoopServiceInstanceListSupplier::new);
        Request finalRequest = request;
        return supplier.get().next().map(serviceInstances -> {
            try {
                return processInstanceResponse(serviceInstances, finalRequest);
            } catch (Exception e) {
                throw new RuntimeException(e);
            }
        });

    }


    private Response<ServiceInstance> processInstanceResponse(List<ServiceInstance> instances, Request request) throws Exception {
        if (instances.isEmpty()) {
            log.warn("No servers available for service: " + this.serviceId);
            return new EmptyResponse();
        } else {
            List<ServiceInstance> serviceInstances = new ArrayList<>();
            HttpHeaders httpHeaders = TransmittableThreadLocalRequestHeaders.get();
            String reqVersion = null;
            if (httpHeaders == null) {
                if (request.getContext() != null) {
                    RequestDataContext requestDataContext = (RequestDataContext) request.getContext();
                    if (requestDataContext != null) {
                        RequestData clientRequest = requestDataContext.getClientRequest();
                        httpHeaders = clientRequest.getHeaders();
                    }
                }
            }
            if (httpHeaders != null) {
                reqVersion = httpHeaders.getFirst(ComConst.VERSION_GRAY_HEADER_KEY);
                log.info("request header version : {}", reqVersion);
            }

            if (StringUtils.isNotBlank(reqVersion)) {
                final String reqVersion1 = reqVersion;
                //判断元数据里是否有version
                serviceInstances = instances.stream()
                        .filter(instance -> instance.getMetadata().get(ComConst.VERSION_GRAY_HEADER_KEY) != null && Objects.equals(instance.getMetadata().get(ComConst.VERSION_GRAY_HEADER_KEY), reqVersion1))
                        .collect(Collectors.toList());
            }

            if (serviceInstances.size() > 0) {
                return processRibbonInstanceResponse(serviceInstances);
            } else {
                List<ServiceInstance> serviceInstances1 = instances.stream()
                        .filter(instance -> instance.getMetadata().get(ComConst.VERSION_GRAY_HEADER_KEY) == null)
                        .collect(Collectors.toList());
                if(serviceInstances1==null||serviceInstances1.size()==0) {
                    throw new Exception("负载均衡找不到服务实例");
                }
                return processRibbonInstanceResponse(serviceInstances1);
            }
        }
    }

    /**
     * 负载均衡器
     * 参考 org.springframework.cloud.loadbalancer.core.RoundRobinLoadBalancer#getInstanceResponse
     *
     * @author javadaily
     */
    private Response<ServiceInstance> processRibbonInstanceResponse(List<ServiceInstance> instances) {
        if (instances == null || instances.size() == 0) {
            return null;
        }
        int pos = Math.abs(this.position.incrementAndGet());
        ServiceInstance instance = instances.get(pos % instances.size());
        return new DefaultResponse(instance);
    }
}